{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  解读 `tvm/script/parser/core/diagnostics.py` "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 解析 `Source`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytest\n",
    "import inspect\n",
    "import tvm.testing\n",
    "from tvm.script.parser.core.diagnostics import Source\n",
    "from tvm.script.parser.core import doc_core as doc\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mInit signature:\u001b[0m \u001b[0mSource\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprogram\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscript\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparser\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdoc_core\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAST\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mDocstring:\u001b[0m     \n",
      "Source code class for TVMScript.\n",
      "\n",
      "It is constructed by source code str or doc AST tree.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "source_name : str\n",
      "    The filename of the file where the source code locates.\n",
      "\n",
      "start_line : int\n",
      "    The first line number of the source code.\n",
      "\n",
      "start_column : int\n",
      "    The first column number of the first line of the source code.\n",
      "\n",
      "source : str\n",
      "    The source code str of source code.\n",
      "\n",
      "full_source : str\n",
      "    The complete source code of the file where the source code locates.\n",
      "\u001b[0;31mFile:\u001b[0m           /media/pc/data/lxw/ai/tvm/python/tvm/script/parser/core/diagnostics.py\n",
      "\u001b[0;31mType:\u001b[0m           type\n",
      "\u001b[0;31mSubclasses:\u001b[0m     "
     ]
    }
   ],
   "source": [
    "Source?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "<module '__main__'> is a built-in module",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 20\u001b[0m\n\u001b[1;32m     12\u001b[0m     source \u001b[38;5;241m=\u001b[39m Source(matmul)\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m (\n\u001b[1;32m     14\u001b[0m         source\u001b[38;5;241m.\u001b[39msource_name \u001b[38;5;241m==\u001b[39m inspect\u001b[38;5;241m.\u001b[39mgetsourcefile(matmul)\n\u001b[1;32m     15\u001b[0m         \u001b[38;5;129;01mand\u001b[39;00m source\u001b[38;5;241m.\u001b[39mstart_line \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[38;5;129;01mand\u001b[39;00m source\u001b[38;5;241m.\u001b[39mfull_source \u001b[38;5;241m==\u001b[39m inspect\u001b[38;5;241m.\u001b[39mgetsource(inspect\u001b[38;5;241m.\u001b[39mgetmodule(matmul))\n\u001b[1;32m     19\u001b[0m     )\n\u001b[0;32m---> 20\u001b[0m \u001b[43mtest_source_base\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[4], line 18\u001b[0m, in \u001b[0;36mtest_source_base\u001b[0;34m()\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtest_source_base\u001b[39m():\n\u001b[1;32m     12\u001b[0m     source \u001b[38;5;241m=\u001b[39m Source(matmul)\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m (\n\u001b[1;32m     14\u001b[0m         source\u001b[38;5;241m.\u001b[39msource_name \u001b[38;5;241m==\u001b[39m inspect\u001b[38;5;241m.\u001b[39mgetsourcefile(matmul)\n\u001b[1;32m     15\u001b[0m         \u001b[38;5;129;01mand\u001b[39;00m source\u001b[38;5;241m.\u001b[39mstart_line \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     16\u001b[0m         \u001b[38;5;129;01mand\u001b[39;00m source\u001b[38;5;241m.\u001b[39mstart_column \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     17\u001b[0m         \u001b[38;5;129;01mand\u001b[39;00m source\u001b[38;5;241m.\u001b[39msource \u001b[38;5;241m==\u001b[39m inspect\u001b[38;5;241m.\u001b[39mgetsource(matmul)\n\u001b[0;32m---> 18\u001b[0m         \u001b[38;5;129;01mand\u001b[39;00m source\u001b[38;5;241m.\u001b[39mfull_source \u001b[38;5;241m==\u001b[39m \u001b[43minspect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgetsource\u001b[49m\u001b[43m(\u001b[49m\u001b[43minspect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgetmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmatmul\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     19\u001b[0m     )\n",
      "File \u001b[0;32m/media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/inspect.py:1282\u001b[0m, in \u001b[0;36mgetsource\u001b[0;34m(object)\u001b[0m\n\u001b[1;32m   1276\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mgetsource\u001b[39m(\u001b[38;5;28mobject\u001b[39m):\n\u001b[1;32m   1277\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Return the text of the source code for an object.\u001b[39;00m\n\u001b[1;32m   1278\u001b[0m \n\u001b[1;32m   1279\u001b[0m \u001b[38;5;124;03m    The argument may be a module, class, method, function, traceback, frame,\u001b[39;00m\n\u001b[1;32m   1280\u001b[0m \u001b[38;5;124;03m    or code object.  The source code is returned as a single string.  An\u001b[39;00m\n\u001b[1;32m   1281\u001b[0m \u001b[38;5;124;03m    OSError is raised if the source code cannot be retrieved.\"\"\"\u001b[39;00m\n\u001b[0;32m-> 1282\u001b[0m     lines, lnum \u001b[38;5;241m=\u001b[39m \u001b[43mgetsourcelines\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mobject\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1283\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(lines)\n",
      "File \u001b[0;32m/media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/inspect.py:1264\u001b[0m, in \u001b[0;36mgetsourcelines\u001b[0;34m(object)\u001b[0m\n\u001b[1;32m   1256\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Return a list of source lines and starting line number for an object.\u001b[39;00m\n\u001b[1;32m   1257\u001b[0m \n\u001b[1;32m   1258\u001b[0m \u001b[38;5;124;03mThe argument may be a module, class, method, function, traceback, frame,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1261\u001b[0m \u001b[38;5;124;03moriginal source file the first line of code was found.  An OSError is\u001b[39;00m\n\u001b[1;32m   1262\u001b[0m \u001b[38;5;124;03mraised if the source code cannot be retrieved.\"\"\"\u001b[39;00m\n\u001b[1;32m   1263\u001b[0m \u001b[38;5;28mobject\u001b[39m \u001b[38;5;241m=\u001b[39m unwrap(\u001b[38;5;28mobject\u001b[39m)\n\u001b[0;32m-> 1264\u001b[0m lines, lnum \u001b[38;5;241m=\u001b[39m \u001b[43mfindsource\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mobject\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1266\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m istraceback(\u001b[38;5;28mobject\u001b[39m):\n\u001b[1;32m   1267\u001b[0m     \u001b[38;5;28mobject\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mobject\u001b[39m\u001b[38;5;241m.\u001b[39mtb_frame\n",
      "File \u001b[0;32m/media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/inspect.py:1075\u001b[0m, in \u001b[0;36mfindsource\u001b[0;34m(object)\u001b[0m\n\u001b[1;32m   1067\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfindsource\u001b[39m(\u001b[38;5;28mobject\u001b[39m):\n\u001b[1;32m   1068\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Return the entire source file and starting line number for an object.\u001b[39;00m\n\u001b[1;32m   1069\u001b[0m \n\u001b[1;32m   1070\u001b[0m \u001b[38;5;124;03m    The argument may be a module, class, method, function, traceback, frame,\u001b[39;00m\n\u001b[1;32m   1071\u001b[0m \u001b[38;5;124;03m    or code object.  The source code is returned as a list of all the lines\u001b[39;00m\n\u001b[1;32m   1072\u001b[0m \u001b[38;5;124;03m    in the file and the line number indexes a line in that list.  An OSError\u001b[39;00m\n\u001b[1;32m   1073\u001b[0m \u001b[38;5;124;03m    is raised if the source code cannot be retrieved.\"\"\"\u001b[39;00m\n\u001b[0;32m-> 1075\u001b[0m     file \u001b[38;5;241m=\u001b[39m \u001b[43mgetsourcefile\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mobject\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1076\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m file:\n\u001b[1;32m   1077\u001b[0m         \u001b[38;5;66;03m# Invalidate cache if needed.\u001b[39;00m\n\u001b[1;32m   1078\u001b[0m         linecache\u001b[38;5;241m.\u001b[39mcheckcache(file)\n",
      "File \u001b[0;32m/media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/inspect.py:952\u001b[0m, in \u001b[0;36mgetsourcefile\u001b[0;34m(object)\u001b[0m\n\u001b[1;32m    948\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mgetsourcefile\u001b[39m(\u001b[38;5;28mobject\u001b[39m):\n\u001b[1;32m    949\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Return the filename that can be used to locate an object's source.\u001b[39;00m\n\u001b[1;32m    950\u001b[0m \u001b[38;5;124;03m    Return None if no way can be identified to get the source.\u001b[39;00m\n\u001b[1;32m    951\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 952\u001b[0m     filename \u001b[38;5;241m=\u001b[39m \u001b[43mgetfile\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mobject\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    953\u001b[0m     all_bytecode_suffixes \u001b[38;5;241m=\u001b[39m importlib\u001b[38;5;241m.\u001b[39mmachinery\u001b[38;5;241m.\u001b[39mDEBUG_BYTECODE_SUFFIXES[:]\n\u001b[1;32m    954\u001b[0m     all_bytecode_suffixes \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m importlib\u001b[38;5;241m.\u001b[39mmachinery\u001b[38;5;241m.\u001b[39mOPTIMIZED_BYTECODE_SUFFIXES[:]\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/script/parser/core/diagnostics.py:110\u001b[0m, in \u001b[0;36m_patched_inspect_getfile\u001b[0;34m(obj)\u001b[0m\n\u001b[1;32m    108\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Work out which source or compiled file an object was defined in.\"\"\"\u001b[39;00m\n\u001b[1;32m    109\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misclass(obj):\n\u001b[0;32m--> 110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_getfile\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobj\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    111\u001b[0m mod \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(obj, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__module__\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m    112\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m mod \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m/media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/inspect.py:913\u001b[0m, in \u001b[0;36mgetfile\u001b[0;34m(object)\u001b[0m\n\u001b[1;32m    911\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mobject\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__file__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m    912\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mobject\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__file__\u001b[39m\n\u001b[0;32m--> 913\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{!r}\u001b[39;00m\u001b[38;5;124m is a built-in module\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\u001b[38;5;28mobject\u001b[39m))\n\u001b[1;32m    914\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m isclass(\u001b[38;5;28mobject\u001b[39m):\n\u001b[1;32m    915\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(\u001b[38;5;28mobject\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__module__\u001b[39m\u001b[38;5;124m'\u001b[39m):\n",
      "\u001b[0;31mTypeError\u001b[0m: <module '__main__'> is a built-in module"
     ]
    }
   ],
   "source": [
    "def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:\n",
    "    A = T.match_buffer(a, [128, 128])\n",
    "    B = T.match_buffer(b, [128, 128])\n",
    "    C = T.match_buffer(c, [128, 128])\n",
    "    for i, j, k in T.grid(128, 128, 128):\n",
    "        with T.block(\"update\"):\n",
    "            vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]\n",
    "\n",
    "\n",
    "def test_source_base():\n",
    "    source = Source(matmul)\n",
    "    assert (\n",
    "        source.source_name == inspect.getsourcefile(matmul)\n",
    "        and source.start_line is not None\n",
    "        and source.start_column == 0\n",
    "        and source.source == inspect.getsource(matmul)\n",
    "        and source.full_source == inspect.getsource(inspect.getmodule(matmul))\n",
    "    )\n",
    "test_source_base()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def test_source_ast():\n",
    "    source = Source(matmul)\n",
    "    mod = source.as_ast()\n",
    "    assert isinstance(mod, doc.Module)\n",
    "    func_def = mod.body[0]\n",
    "    assert isinstance(func_def, doc.FunctionDef)\n",
    "    assert func_def.name == \"matmul\"\n",
    "    func_args = func_def.args\n",
    "    assert (\n",
    "        len(func_args.args) == 3\n",
    "        and func_args.args[0].arg == \"a\"\n",
    "        and func_args.args[1].arg == \"b\"\n",
    "        and func_args.args[2].arg == \"c\"\n",
    "    )\n",
    "    func_body = func_def.body\n",
    "    assert len(func_body) == 4\n",
    "    func_assigns = func_body[:3]\n",
    "    assert (\n",
    "        isinstance(func_assigns[0], doc.Assign)\n",
    "        and func_assigns[0].targets[0].id == \"A\"\n",
    "        and isinstance(func_assigns[1], doc.Assign)\n",
    "        and func_assigns[1].targets[0].id == \"B\"\n",
    "        and isinstance(func_assigns[2], doc.Assign)\n",
    "        and func_assigns[2].targets[0].id == \"C\"\n",
    "    )\n",
    "    func_for = func_body[3]\n",
    "    assert (\n",
    "        len(func_for.target.elts) == 3\n",
    "        and func_for.target.elts[0].id == \"i\"\n",
    "        and func_for.target.elts[1].id == \"j\"\n",
    "        and func_for.target.elts[2].id == \"k\"\n",
    "    )\n",
    "    for_body = func_for.body\n",
    "    assert len(for_body) == 1\n",
    "    for_block = for_body[0]\n",
    "    assert isinstance(for_block, doc.With) and len(for_block.body) == 2\n",
    "\n",
    "\n",
    "def test_nesting_parsing():\n",
    "    class dummy:\n",
    "        pass\n",
    "\n",
    "    for i in range(1):\n",
    "\n",
    "        @tvm.script.ir_module\n",
    "        class Module:\n",
    "            @T.prim_func\n",
    "            def impl(\n",
    "                A: T.Buffer((12, 196, 64), \"float32\"),\n",
    "            ) -> None:\n",
    "                T.evaluate(0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
